package gorm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"strings"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

var (
	writeStmt  = regexp.MustCompile(`(?i)^\s*(INSERT|UPDATE|DELETE)`)
	selectStmt = regexp.MustCompile(`(?i)^\s*(SELECT|DESCRIBE|SHOW|DESC)`)
	createStmt = regexp.MustCompile(`(?i)^\s*CREATE TABLE`)
)

// HandleQueryTool handles SQL query execution
func HandleQueryTool(dbClient *DBClient) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		arguments := request.Params.Arguments
		message, ok := arguments["sql"].(string)
		if !ok {
			return nil, fmt.Errorf("invalid message argument")
		}

		results, err := dbClient.ExecuteSQL(message)
		if err != nil {
			return nil, fmt.Errorf("failed to execute SQL query: %w", err)
		}

		jsonData, err := json.Marshal(results)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal SQL results: %w", err)
		}

		return &mcp.CallToolResult{
			Content: []mcp.Content{
				mcp.TextContent{
					Type: "text",
					Text: string(jsonData),
				},
			},
		}, nil
	}
}

func ReadQueryToolHandler(dbClient *DBClient) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		arguments := request.Params.Arguments
		message, ok := arguments["query"].(string)
		if !ok {
			return nil, fmt.Errorf("invalid message argument")
		}

		if !selectStmt.MatchString(message) {
			return nil, errors.New("only SELECT queries are allowed")
		}

		results, err := dbClient.ExecuteSQL(message)
		if err != nil {
			return nil, fmt.Errorf("failed to execute SQL query: %w", err)
		}

		jsonData, err := json.Marshal(results)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal SQL results: %w", err)
		}

		return &mcp.CallToolResult{
			Content: []mcp.Content{
				mcp.TextContent{
					Type: "text",
					Text: string(jsonData),
				},
			},
		}, nil
	}
}

func WriteQueryToolHandler(dbClient *DBClient) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		arguments := request.Params.Arguments
		message, ok := arguments["query"].(string)
		if !ok {
			return nil, errors.New("invalid query parameter")
		}

		if !writeStmt.MatchString(message) {
			return nil, errors.New("only INSERT/UPDATE/DELETE queries are allowed")
		}

		results, err := dbClient.ExecuteSQL(message)
		if err != nil {
			return nil, fmt.Errorf("Write operation failed: %w", err)
		}

		jsonData, err := json.Marshal(results)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal SQL results: %w", err)
		}
		return &mcp.CallToolResult{
			Content: []mcp.Content{
				mcp.TextContent{
					Type: "text",
					Text: string(jsonData),
				},
			},
		}, nil
	}
}

func CreateTableToolHandler(dbClient *DBClient) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		arguments := request.Params.Arguments
		schema, ok := arguments["schema"].(string)
		if !ok {
			return nil, errors.New("invalid schema parameter")
		}

		if !createStmt.MatchString(schema) {
			return nil, errors.New("invalid CREATE TABLE statement")
		}

		results, err := dbClient.ExecuteSQL(schema)
		if err != nil {
			return nil, fmt.Errorf("Write operation failed: %w", err)
		}

		jsonData, err := json.Marshal(results)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal SQL results: %w", err)
		}

		return &mcp.CallToolResult{
			Content: []mcp.Content{
				mcp.TextContent{
					Type: "text",
					Text: string(jsonData),
				},
			},
		}, nil
	}
}

func sanitizeInput(input string) string {
	return strings.ReplaceAll(input, "'", "''")
}

func ListTableToolHandler(dbClient *DBClient) server.ToolHandlerFunc {
	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		var schemaFilter string
		dbType := dbClient.db.Dialector.Name()
		if schema, ok := request.Params.Arguments["schema"].(string); ok {

			switch dbType {
			case "mysql":
				schemaFilter = fmt.Sprintf(" AND table_schema = '%s'", sanitizeInput(schema))
			case "postgres":
				schemaFilter = fmt.Sprintf(" AND schemaname = '%s'", sanitizeInput(schema))
			case "clickhouse":
				schemaFilter = fmt.Sprintf(" AND database = '%s'", sanitizeInput(schema))
				// SQLite 没有 schema 概念，所以不需要 schemaFilter
			}
		}

		query := ""
		if dbType == "mysql" {
			query = fmt.Sprintf(`
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys') %s`, schemaFilter)
		} else if dbType == "postgres" {
			query = fmt.Sprintf(`
        SELECT tablename 
        FROM pg_catalog.pg_tables 
        WHERE schemaname NOT IN ('pg_catalog', 'information_schema') %s
    `, schemaFilter)
		} else if dbType == "sqlite" {
			query = `
        SELECT name as table_name
        FROM sqlite_master
        WHERE type='table' AND name NOT LIKE 'sqlite_%'`
		} else if dbType == "clickhouse" {
			query = fmt.Sprintf(`
        SELECT name as table_name
        FROM system.tables
        WHERE database NOT IN ('system', 'information_schema') %s
    `, schemaFilter)
		} else {
			return nil, fmt.Errorf("unsupported database type: %s", dbType)
		}

		results, err := dbClient.ExecuteSQL(query)
		if err != nil {
			return nil, fmt.Errorf("Write operation failed: %w", err)
		}

		jsonData, err := json.Marshal(results)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal SQL results: %w", err)
		}

		return &mcp.CallToolResult{
			Content: []mcp.Content{
				mcp.TextContent{
					Type: "text",
					Text: string(jsonData),
				},
			},
		}, nil

	}

}

// GetQueryToolSchema returns the schema for query tool
func GetQueryToolSchema() json.RawMessage {
	return json.RawMessage(`
	{
		"type": "object",
		"properties": {
		"sql": { 
				"type": "string",
				"description": "The sql query to execute"
			}
		}
	}
	`)
}
